package com.lzj.server.common.service.impl;

import com.alibaba.dashscope.embeddings.TextEmbedding;
import com.alibaba.dashscope.embeddings.TextEmbeddingParam;
import com.alibaba.dashscope.embeddings.TextEmbeddingResult;
import com.alibaba.dashscope.embeddings.TextEmbeddingResultItem;
import com.alibaba.dashscope.exception.ApiException;
import com.alibaba.dashscope.exception.NoApiKeyException;
import com.lzj.server.common.model.VideoVectorModel;
import com.lzj.server.common.service.DashScopeEmbeddingService;
import com.lzj.server.entity.video.VideoLibrary;
import com.lzj.server.exceptions.BaseException;
import com.lzj.server.repository.mapper.VideoVectorModelMapper;
import com.lzj.server.service.VideoService;
import jakarta.annotation.Resource;
import lombok.extern.slf4j.Slf4j;
import org.springframework.stereotype.Service;

import java.util.List;

/**
 * DashScope Embedding 调用服务
 */
@Slf4j
@Service
public class DashScopeEmbeddingServiceImpl implements DashScopeEmbeddingService {
    @Resource
    private VideoVectorModelMapper videoVectorModelMapper;
    @Resource
    private VideoService videoLibraryService;

    private static final int DASHSCOPE_MAX_BATCH_SIZE = 25;

    @Override
    public Long saveEmbeddings(Long id) {
        // 查询视频库
        VideoLibrary videoLibrary = videoLibraryService.getById(id);
        if (videoLibrary == null) {
            throw new BaseException("视频库不存在");
        }
        // 调用通义千问 embedding
        String keyword = "标题：" + videoLibrary.getTitle() + "，标签：" + videoLibrary.getVideoLabel();
        List<String> inputs = List.of(keyword);

        TextEmbeddingResult result = null;
        int batchCounter = 0;

        for (int i = 0; i < inputs.size(); i += DASHSCOPE_MAX_BATCH_SIZE) {
            List<String> batch = inputs.subList(i, Math.min(i + DASHSCOPE_MAX_BATCH_SIZE, inputs.size()));
            TextEmbeddingParam param = TextEmbeddingParam.builder()
                    .model("text-embedding-v4")
                    .texts(batch)
                    .parameter("dimension", 1024)  // 指定向量维度（仅 text-embedding-v3及 text-embedding-v4支持该参数）
                    .build();

            TextEmbedding textEmbedding = new TextEmbedding();
            try {
                TextEmbeddingResult resp = textEmbedding.call(param);
                if (resp != null) {
                    if (result == null) {
                        result = resp;
                    } else {
                        for (var emb : resp.getOutput().getEmbeddings()) {
                            emb.setTextIndex(emb.getTextIndex() + batchCounter);
                            result.getOutput().getEmbeddings().add(emb);
                        }
                        result.getUsage().setTotalTokens(result.getUsage().getTotalTokens() + resp.getUsage().getTotalTokens());
                    }
                } else {
                    throw new BaseException("Embedding调用失败");
                }
            } catch (ApiException | NoApiKeyException e) {
                throw new BaseException("Embedding调用失败", e);
            }
            batchCounter += batch.size();
        }
        List<Double> embedding = result.getOutput().getEmbeddings().get(0).getEmbedding();
        // 封装视频向量数据
        VideoVectorModel v = new VideoVectorModel();
        v.setVideoId(videoLibrary.getId());
        v.setAttachmentId(videoLibrary.getVideoAttachmentId());
        v.setVideoVector(embedding.stream().mapToDouble(Double::doubleValue).toArray());
        v.setAttachmentVector(embedding.stream().mapToDouble(Double::doubleValue).toArray());
        // 保存到数据库
        videoVectorModelMapper.insert(v);
        System.out.println(result);
        return v.getId();
    }

    @Override
    public TextEmbeddingResult getEmbeddings(List<String> texts, String model, int dimension) {
        try {
            TextEmbeddingParam param = TextEmbeddingParam.builder()
                    .model(model)
                    .texts(texts)
                    .parameter("dimension", 1024)
                    .build();

            TextEmbedding textEmbedding = new TextEmbedding();
            return textEmbedding.call(param);

        } catch (ApiException | NoApiKeyException e) {
            log.error("Embedding 失败: {}", e.getMessage(), e);
            throw new BaseException("Embedding调用失败", e);
        }
    }

    public TextEmbeddingResult getBatchEmbeddings(List<String> inputs, int batchSize, String model, int dimension) {
        TextEmbeddingResult finalResult = null;
        int totalIndexOffset = 0;

        for (int i = 0; i < inputs.size(); i += batchSize) {
            List<String> batch = inputs.subList(i, Math.min(i + batchSize, inputs.size()));

            TextEmbeddingParam param = TextEmbeddingParam.builder()
                    .model(model)
                    .texts(batch)
                    .parameter("dimension", dimension)
                    .build();

            TextEmbedding embeddingClient = new TextEmbedding();
            try {
                TextEmbeddingResult resp = embeddingClient.call(param);
                if (resp != null) {
                    List<TextEmbeddingResultItem> embeddings = resp.getOutput().getEmbeddings();
                    // 手动补正 index
                    for (int j = 0; j < embeddings.size(); j++) {
                        embeddings.get(j).setTextIndex(totalIndexOffset + j);
                    }
                    // 合并结果
                    if (finalResult == null) {
                        finalResult = resp;
                    } else {
                        finalResult.getOutput().getEmbeddings().addAll(embeddings);
                        finalResult.getUsage().setTotalTokens(
                                finalResult.getUsage().getTotalTokens() + resp.getUsage().getTotalTokens()
                        );
                    }
                }
            } catch (ApiException | NoApiKeyException e) {
                throw new BaseException("LLM调用失败", e);
            }
            totalIndexOffset += batch.size();
        }
        return finalResult;
    }
}
